#include "compi/context.h"
#include "compi/fixed_type.h"
#include "compi/resource_manager.h"

#include <atomic>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <mutex>
#include <stdexcept>
#include <unordered_map>

namespace compi
{

// 全局失败策略：当无法获得 Context/Environment 时用于决定失败行为
static std::atomic<FailurePolicy> g_global_failure_policy{FailurePolicy::ThrowException};

FailurePolicy
global_failure_policy() noexcept
{
    return g_global_failure_policy.load();
}

void
set_global_failure_policy(FailurePolicy policy) noexcept
{
    g_global_failure_policy.store(policy);
}

// 本地哈希函数：避免依赖 ResourceManager 的私有 compute_comm_hash
static std::size_t
hash_comm(MPI_Comm comm)
{
    MPI_Fint handle = MPI_Comm_c2f(comm);
    std::hash<MPI_Fint> hasher;
    return hasher(handle);
}

// Context 的 pimpl 实现
struct Context::Impl
{
    MPI_Comm comm;
    int rank = -1;
    int size = -1;
    MpiStats stats;
    std::unique_ptr<ErrorStack> error_stack;
    Environment * env = nullptr;
    MessageTagAllocator * tag_allocator = nullptr;
    bool has_failure_override = false;
    FailurePolicy failure_override = FailurePolicy::ThrowException;
    std::mutex stats_mutex;

    explicit Impl(MPI_Comm c) : comm(c), error_stack(std::make_unique<ErrorStack>()) {}
};

// 全局单例注册表（使用稳定哈希键）
static std::mutex g_context_mutex;
static std::unordered_map<std::size_t, std::unique_ptr<Context>> g_contexts;

Context &
Context::for_comm(MPI_Comm comm)
{
    std::lock_guard<std::mutex> lock(g_context_mutex);

    std::size_t hash = hash_comm(comm);
    auto it = g_contexts.find(hash);
    if (it != g_contexts.end())
    {
        // 验证 communicator 是否匹配，否则移除冲突条目
        int cmp = MPI_UNEQUAL;
        (void)MPI_Comm_compare(comm, it->second->comm(), &cmp);
        if (cmp == MPI_IDENT || cmp == MPI_CONGRUENT)
        {
            return *it->second;
        }
        // 哈希冲突，移除旧条目
        g_contexts.erase(it);
    }

    // 创建新的 Context 实例
    auto context = std::make_unique<Context>(comm);
    auto * ptr = context.get();
    g_contexts[hash] = std::move(context);

    return *ptr;
}

Context &
Context::for_comm(MPI_Comm comm, Environment & env)
{
    Context & ctx = for_comm(comm);
    // 记录环境引用（若尚未设置）
    if (ctx._impl->env == nullptr)
    {
        ctx._impl->env = &env;
    }
    return ctx;
}

void
Context::destroy_for_comm(MPI_Comm comm)
{
    // 统一生命周期：在销毁 Context 之前，释放与该 communicator 关联的资源单例
    // 尤其是标签分配器，避免“真单例”在并发与多入口下产生归属不唯一的问题。
    ResourceManager::destroy_tag_allocator(comm);

    std::lock_guard<std::mutex> lock(g_context_mutex);
    std::size_t hash = hash_comm(comm);
    g_contexts.erase(hash);
}

Context::Context(MPI_Comm comm) : _impl(std::make_unique<Impl>(comm))
{
    initialize_comm_properties();
}

Context::~Context() = default;

bool
Context::initialize_comm_properties()
{
    int rank, size;
    if (MPI_Comm_rank(_impl->comm, &rank) != MPI_SUCCESS)
    {
        return false;
    }
    if (MPI_Comm_size(_impl->comm, &size) != MPI_SUCCESS)
    {
        return false;
    }

    _impl->rank = rank;
    _impl->size = size;
    return true;
}

Result<void>
Context::call_mpi(std::function<int()> mpi_fn,
                  const char * operation_name,
                  std::optional<FailurePolicy> policy_override)
{
    auto start_time = std::chrono::high_resolution_clock::now();

    {
        std::lock_guard<std::mutex> lock(_impl->stats_mutex);
        _impl->stats.total_calls++;
    }

    int mpi_result = mpi_fn();

    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
    {
        std::lock_guard<std::mutex> lock(_impl->stats_mutex);
        _impl->stats.total_time_ms += duration.count() / 1000.0;
    }

    // 统一失败策略处理（由 Context 注入策略）
    if (!compi::call_mpi(*this, mpi_result, operation_name, policy_override))
    {
        {
            std::lock_guard<std::mutex> lock(_impl->stats_mutex);
            _impl->stats.failed_calls++;
        }
        std::string error_msg = std::string("MPI operation failed: ") + operation_name;
        _impl->error_stack->push_error(ResultCode::MpiError, error_msg, operation_name);
        return Result<void>::error(ResultCode::MpiError, error_msg, operation_name);
    }

    return COMPI_SUCCESS_VOID_OP(operation_name);
}

MPI_Comm
Context::comm() const noexcept
{
    return _impl->comm;
}

int
Context::rank() const noexcept
{
    return _impl->rank;
}

int
Context::size() const noexcept
{
    return _impl->size;
}

ErrorStack &
Context::error_stack() noexcept
{
    return *_impl->error_stack;
}

Context::MpiStats
Context::stats() const
{
    std::lock_guard<std::mutex> lock(_impl->stats_mutex);
    return _impl->stats;
}

void
Context::reset_stats() noexcept
{
    std::lock_guard<std::mutex> lock(_impl->stats_mutex);
    _impl->stats = MpiStats{};
}

// 资源访问（来自原 Context）
MessageTagAllocator &
Context::tag_allocator()
{
    if (_impl->tag_allocator == nullptr)
    {
        _impl->tag_allocator = &ResourceManager::get_tag_allocator(_impl->comm);
    }
    return *_impl->tag_allocator;
}

Result<Environment &>
Context::environment() const noexcept
{
    if (_impl->env)
    {
        return Result<Environment &>::success(*_impl->env);
    }
    return Result<Environment &>::error(
        ResultCode::InvalidRequest,
        "Environment is not set. Call Context::for_comm(comm, env) before use.");
}

// 失败策略访问器
void
Context::set_failure_policy(FailurePolicy policy) noexcept
{
    _impl->has_failure_override = true;
    _impl->failure_override = policy;
}

FailurePolicy
Context::failure_policy() const noexcept
{
    if (_impl->has_failure_override)
    {
        return _impl->failure_override;
    }
    if (_impl->env)
    {
        return _impl->env->failure_policy();
    }
    return FailurePolicy::ThrowException;
}

FailurePolicy
Context::effective_failure_policy() const noexcept
{
    return failure_policy();
}

// 基于 Result<void> 的统一错误处理：写入错误栈并按策略执行行为
void
Context::handle_error(const Result<void> & res)
{
    // 记录到错误栈
    if (_impl && _impl->error_stack)
    {
        _impl->error_stack->push_error(res.error_code(), res.error_message(), res.operation());
    }

    const FailurePolicy policy = effective_failure_policy();
    switch (policy)
    {
        case FailurePolicy::ThrowException:
            throw std::runtime_error(res.error_message());
        case FailurePolicy::LogAndContinue:
            std::fprintf(stderr,
                         "[COMPI] Error(%d) in %s: %s\n",
                         static_cast<int>(res.error_code()),
                         res.operation().c_str(),
                         res.error_message().c_str());
            std::fflush(stderr);
            return; // 继续执行
        case FailurePolicy::AbortProcess:
            // 使用状态码作为退出码（不要求是 MPI 错误码）
            MPI_Abort(comm(), static_cast<int>(res.error_code()));
            return;
    }
}

namespace
{
static int
to_mpi_thread_level(ThreadLevel lvl)
{
    switch (lvl)
    {
        case ThreadLevel::Single:
            return MPI_THREAD_SINGLE;
        case ThreadLevel::Funneled:
            return MPI_THREAD_FUNNELED;
        case ThreadLevel::Serialized:
            return MPI_THREAD_SERIALIZED;
        case ThreadLevel::Multiple:
            return MPI_THREAD_MULTIPLE;
    }
    return MPI_THREAD_SINGLE;
}

static ThreadLevel
from_mpi_thread_level(int lvl)
{
    switch (lvl)
    {
        case MPI_THREAD_SINGLE:
            return ThreadLevel::Single;
        case MPI_THREAD_FUNNELED:
            return ThreadLevel::Funneled;
        case MPI_THREAD_SERIALIZED:
            return ThreadLevel::Serialized;
        case MPI_THREAD_MULTIPLE:
            return ThreadLevel::Multiple;
        default:
            return ThreadLevel::Single;
    }
}

// ANY_TAG 命名空间过滤辅助：在非阻塞匹配中规避队首阻塞，扫描命名空间可能的标签。
static int
query_tag_ub_local(MPI_Comm comm)
{
    int flag = 0;
    int * val = nullptr;
    int rc = MPI_Comm_get_attr(comm, MPI_TAG_UB, &val, &flag);
    if (rc == MPI_SUCCESS && flag && val)
        return *val;
    return 32767; // MPI 至少保证该上限
}

static int
compute_stride_local(MPI_Comm comm)
{
    const int tag_ub = query_tag_ub_local(comm);
    int candidate = std::max(64, std::min(1024, tag_ub / 8));
    if (candidate <= 0)
        candidate = 64;
    return candidate;
}
}

// 初始化 MPI 环境（线程化）。如果尚未初始化则使用 MPI_Init_thread；否则查询已提供的线程级别。
// 参数：req 请求的线程级别。
// 返回值：无
// 异常：当 MPI 初始化或查询失败时，按照失败策略处理。
Environment::Environment(ThreadLevel req)
{
    int initialized = 0;
    MPI_Initialized(&initialized);
    int provided = MPI_THREAD_SINGLE;
    if (!initialized)
    {
        int argc = 0;
        char ** argv = nullptr;
        const int required = to_mpi_thread_level(req);
        const int rc = MPI_Init_thread(&argc, &argv, required, &provided);
        (void)call_mpi(rc, "MPI_Init_thread failed", MPI_COMM_WORLD);
        // 本实例负责调用 MPI_Init_*，在析构中执行 MPI_Finalize
        _owns_mpi = true;
    }
    else
    {
        // 已初始化：查询线程级别
        const int rc = MPI_Query_thread(&provided);
        (void)call_mpi(rc, "MPI_Query_thread failed", MPI_COMM_WORLD);
        // 外部已初始化，本实例不负责 finalize
        _owns_mpi = false;
    }
    _requested_level = req;
    _provided_level = from_mpi_thread_level(provided);
}

// 终止 MPI 环境。如果尚未终止则进行终止；否则不做处理。
// 参数：无
// 返回值：无
// 异常：当 MPI_Finalize 返回非成功码时，按照失败策略处理。
Environment::~Environment()
{
    // 仅当由本实例负责初始化时才执行最终化，避免重复 MPI_Finalize
    if (_owns_mpi)
    {
        int finalized = 0;
        MPI_Finalized(&finalized);
        if (!finalized)
        {
            // 在 MPI_Finalize 之前释放派生 MPI_Datatype，避免 finalize 后释放导致未定义行为
            compi::types::cleanup_complex_types_before_finalize();
            const int rc = MPI_Finalize();
            // 注意：MPI 已处于终止过程，禁止再通过 Context 封装调用任何 MPI 例程。
            // 这里改为直接检查返回码并记录日志，避免在 finalize 之后触发 MPI_Comm_rank 等调用。
            if (rc != MPI_SUCCESS)
            {
                std::fprintf(stderr, "[COMPI] MPI_Finalize failed (rc=%d)\n", rc);
                std::fflush(stderr);
            }
        }
    }
}

void
Environment::set_failure_policy(FailurePolicy policy) noexcept
{
    _failure_policy = policy;
    // 同步更新全局失败策略，供无法获得 Context 的路径使用
    set_global_failure_policy(policy);
}

FailurePolicy
Environment::failure_policy() const noexcept
{
    return _failure_policy;
}

// ErrorStack的pimpl实现
struct ErrorStack::Impl
{
    std::vector<ErrorEntry> error_history;
    std::size_t max_history_size = 100;
    std::mutex mutex;
};

ErrorStack::ErrorStack() : _impl(std::make_unique<Impl>()) {}

ErrorStack::~ErrorStack() = default;

ErrorStack::ErrorStack(ErrorStack && other) noexcept : _impl(std::move(other._impl)) {}

ErrorStack &
ErrorStack::operator=(ErrorStack && other) noexcept
{
    if (this != &other)
    {
        _impl = std::move(other._impl);
    }
    return *this;
}

void
ErrorStack::push_error(ResultCode code, const std::string & message, const std::string & operation)
{
    std::lock_guard<std::mutex> lock(_impl->mutex);
    ErrorEntry entry{.code = code,
                     .message = message,
                     .operation = operation,
                     .timestamp = std::chrono::system_clock::now()};

    _impl->error_history.push_back(std::move(entry));

    // 保持错误历史大小限制
    if (_impl->error_history.size() > _impl->max_history_size)
    {
        _impl->error_history.erase(_impl->error_history.begin());
    }
}

Result<ErrorStack::ErrorEntry>
ErrorStack::get_last_error() const
{
    std::lock_guard<std::mutex> lock(_impl->mutex);
    if (_impl->error_history.empty())
    {
        return Result<ErrorEntry>::error(ResultCode::Unknown, "No errors in history");
    }
    return Result<ErrorEntry>::success(_impl->error_history.back());
}

const ErrorStack::ErrorEntry *
ErrorStack::last_error() const noexcept
{
    std::lock_guard<std::mutex> lock(_impl->mutex);
    if (_impl->error_history.empty())
    {
        return nullptr;
    }
    return &_impl->error_history.back();
}

std::vector<ErrorStack::ErrorEntry>
ErrorStack::all_errors() const
{
    std::lock_guard<std::mutex> lock(_impl->mutex);
    return _impl->error_history;
}

void
ErrorStack::clear() noexcept
{
    std::lock_guard<std::mutex> lock(_impl->mutex);
    _impl->error_history.clear();
}

bool
ErrorStack::has_errors() const noexcept
{
    std::lock_guard<std::mutex> lock(_impl->mutex);
    return !_impl->error_history.empty();
}

std::size_t
ErrorStack::error_count() const noexcept
{
    std::lock_guard<std::mutex> lock(_impl->mutex);
    return _impl->error_history.size();
}

} // namespace compi

namespace compi
{

// 统一 MPI 返回码处理：依赖 Context 的策略
bool
call_mpi(Context & ctx,
         int rc,
         const char * operation_name,
         std::optional<FailurePolicy> policy_override)
{
    if (rc == MPI_SUCCESS)
        return true;

    const FailurePolicy policy =
        policy_override.has_value() ? *policy_override : ctx.effective_failure_policy();
    switch (policy)
    {
        case FailurePolicy::ThrowException:
            throw std::runtime_error(std::string("MPI call failed: ") + operation_name);
        case FailurePolicy::LogAndContinue:
            std::fprintf(stderr, "[COMPI] MPI error: %s\n", operation_name);
            std::fflush(stderr);
            return false;
        case FailurePolicy::AbortProcess:
            MPI_Abort(ctx.comm(), rc);
            return false;
    }
    // 防御性返回
    return false;
}

} // namespace compi
